# --------------------------------------------------------
# References:
# https://github.com/jxhe/unify-parameter-efficient-tuning
# --------------------------------------------------------

import math
import torch
import torch.nn as nn
from utils.toolkit import NamespaceDict

# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# MAE: https://github.com/facebookresearch/mae
# --------------------------------------------------------
from timm.models.vision_transformer import PatchEmbed
from timm.models.registry import register_model
from typing import Optional

from models.components.backbone.utils import initialize_vit_model
from models.components.backbone.vit import VisionTransformer as _VisionTransformer


class Lora(nn.Module):
    def __init__(
        self,
        tuning_config: NamespaceDict,
        embed_dim: Optional[int] = None,
        bottleneck: Optional[int] = None,
        dropout: float = 0.0,
        adapter_scalar="1.0",
        adapter_layernorm_option="in",
    ):
        super().__init__()
        self._device = tuning_config._device
        self.n_embd = tuning_config.embed_dim if embed_dim is None else embed_dim
        self.down_size = tuning_config.attn_bn if bottleneck is None else bottleneck

        # layernorm
        self.adapter_layernorm_option = adapter_layernorm_option

        if adapter_layernorm_option == "in" or adapter_layernorm_option == "out":
            self.adapter_layer_norm = nn.LayerNorm(self.n_embd)

        if adapter_scalar == "learnable_scalar":
            self.scale = nn.Parameter(torch.ones(1))
        else:
            self.scale = float(adapter_scalar)

        self.down_proj = nn.Linear(self.n_embd, self.down_size, bias=False)
        self.act_func = nn.Identity()
        self.up_proj = nn.Linear(self.down_size, self.n_embd, bias=False)

        self.dropout = dropout

        with torch.no_grad():
            nn.init.kaiming_uniform_(self.down_proj.weight, a=math.sqrt(5))
            nn.init.zeros_(self.up_proj.weight)

    def forward(self, x, add_residual=True, residual=None):
        return self._forward(x, add_residual, residual)

    def _update(self):
        print("Update Lora: Not Implemented")

    def _forward(self, x, add_residual=True, residual=None):
        residual = x if residual is None else residual
        if self.adapter_layernorm_option == "in":
            x = self.adapter_layer_norm(x)

        down = self.down_proj(x)
        down = self.act_func(down)
        down = nn.functional.dropout(down, p=self.dropout, training=self.training)
        up = self.up_proj(down)

        up = up * self.scale

        if self.adapter_layernorm_option == "out":
            up = self.adapter_layer_norm(up)

        if add_residual:
            output = up + residual
        else:
            output = up

        return output


class VisionTransformer(_VisionTransformer):
    """Vision Transformer with support for global average pooling"""

    def __init__(
        self,
        global_pool=False,
        img_size=224,
        patch_size=16,
        in_chans=3,
        num_classes=1000,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        qkv_bias=True,
        representation_size=None,
        distilled=False,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.0,
        embed_layer=PatchEmbed,
        norm_layer=None,
        act_layer=None,
        weight_init="",
        config: NamespaceDict = NamespaceDict(),
    ):
        super(VisionTransformer, self).__init__(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            num_classes=num_classes,
            embed_dim=embed_dim,
            depth=depth,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            representation_size=representation_size,
            distilled=distilled,
            drop_rate=drop_rate,
            attn_drop_rate=attn_drop_rate,
            drop_path_rate=drop_path_rate,
            embed_layer=embed_layer,
            norm_layer=norm_layer,
            act_layer=act_layer,
            weight_init=weight_init,
            global_pool=global_pool,
            config=config,
        )
        self.cur_adapter = nn.ModuleList()
        self.get_new_adapter()

    def init_weights(self, mode=""):
        raise NotImplementedError()

    @torch.jit.ignore  # type: ignore
    def no_weight_decay(self):
        return {"pos_embed", "cls_token", "dist_token"}

    def get_classifier(self):
        if self.dist_token is None:
            return self.head
        else:
            return self.head, self.head_dist

    def reset_classifier(self, num_classes, global_pool=""):
        self.num_classes = num_classes
        self.head = (
            nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
        )
        if self.num_tokens == 2:
            self.head_dist = (
                nn.Linear(self.embed_dim, self.num_classes)
                if num_classes > 0
                else nn.Identity()
            )

    def freeze(self):
        for param in self.parameters():
            param.requires_grad = False

        for i in range(len(self.cur_adapter)):
            self.cur_adapter[i].requires_grad = True  # type: ignore

    def after_medium(self):
        pass

    def after_task(self):
        pass

    def forward_token(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        for idx, blk in enumerate(self.blocks):
            x = blk(x, self.cur_adapter[idx])

        # whether to use class token or mean of all tokens
        if self.global_pool:
            x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
            outcome = self.fc_norm(x)
        else:
            x = self.norm(x)
            outcome = x[:, 0]

        return outcome

    def forward_feats(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        for idx, blk in enumerate(self.blocks):
            x = blk(x, self.cur_adapter[idx])

        x = self.norm(x)

        return x

    def forward(self, x):
        x = self.forward_token(x)
        return x

    def get_new_adapter(self):
        config = self.config
        self.cur_adapter = nn.ModuleList()
        if config.ffn_adapt:
            for i in range(len(self.blocks)):
                adapter = Lora(
                    self.config,
                    embed_dim=self.embed_dim,
                    bottleneck=config.ffn_rank,
                    dropout=0.1,
                    adapter_scalar=config.ffn_adapter_scalar,
                    adapter_layernorm_option=config.ffn_adapter_layernorm_option,
                ).to(self._device)
                self.cur_adapter.append(adapter)
            self.cur_adapter.requires_grad_(True)
        else:
            print("====Not use adapter===")


@register_model
def vit_base_patch16_224_baseSingle(pretrained=False, pretrained_cfg={}, **kwargs):
    del pretrained
    del pretrained_cfg
    return initialize_vit_model(
        "vit_base_patch16_224",
        VisionTransformer,
        **kwargs,
    )


@register_model
def vit_base_patch16_224_in21k_baseSingle(
    pretrained=False, pretrained_cfg={}, **kwargs
):
    del pretrained
    del pretrained_cfg
    return initialize_vit_model(
        "vit_base_patch16_224_in21k",
        VisionTransformer,
        **kwargs,
    )
